from gymnasium.envs.registration import register

register(
    id='seals/pih-v0',
    entry_point='gym_pih.envs:make_wrapped_env',
    max_episode_steps=200,
)

register(
    id='custom/pih-v0',
    entry_point='gym_pih.envs:TypecEnv',
    max_episode_steps=200,
)

register(
    id='serl/pih-v0',
    entry_point='gym_pih.envs:TypecEnv',
    max_episode_steps=200,
    kwargs={'train_env': 'serl',
            'fake_env':False}
)

register(
    id='serl/pih-v2',
    entry_point='gym_pih.envs:TypecEnv_v2',
    max_episode_steps=200,
    kwargs={'train_env': 'serl',
            'fake_env':False}
)

register(
    id='serl/pih-v3',
    entry_point='gym_pih.envs:TypecEnv_v3',
    max_episode_steps=200,
    kwargs={'train_env': 'serl',
            'fake_env':False}
)

register(
    id='serl/pih-v4',
    entry_point='gym_pih.envs:TypecEnv_v4',
    max_episode_steps=200,
    kwargs={'train_env': 'serl',
            'fake_env':False}
)

register(
    id='serl/pih-v5',
    entry_point='gym_pih.envs:TypecEnv_v5',
    max_episode_steps=200,
    kwargs={'train_env': 'serl',
            'fake_env':False}
)

register(
    id='two-stage',
    entry_point='gym_pih.envs:TwoStageEnv',
    max_episode_steps=20,
)

register(
    id='act',
    entry_point='gym_pih.envs.act:ACTEnv',
    max_episode_steps=400,
)

register(
    id='act_enhance',
    entry_point='gym_pih.envs.act_image_enhancement:ActImageEnhancement',
    max_episode_steps=400,
)